"""
游戏AI代理类
"""
import numpy as np
import random
from ball_battle_env import MultiAgentBallBattleEnv


class RandomAgent:
    """随机动作选择的AI代理"""
    
    def __init__(self, action_space):
        self.action_space = action_space
    
    def select_action(self, observation):
        """根据当前观察选择随机动作"""
        return self.action_space.sample()


class GameAI:
    """游戏AI控制器，管理AI代理与游戏的交互"""
    
    def __init__(self, render_mode=None):
        # 创建环境
        self.env = MultiAgentBallBattleEnv(render_mode=render_mode)
        
        # 创建AI代理（初始为随机代理）
        self.agent = RandomAgent(self.env.action_space)
        
        # 记录游戏状态
        self.current_observation = None
        self.current_info = None
        self.episode_reward = 0
        self.episode_count = 0
        
    def reset(self):
        """重置游戏和环境"""
        self.current_observation, self.current_info = self.env.reset()
        self.episode_reward = 0
        return self.current_observation
    
    def step(self):
        """执行一步游戏"""
        if self.current_observation is None:
            self.reset()
        
        # AI代理选择动作
        action = self.agent.select_action(self.current_observation)
        
        # 执行动作
        observation, reward, terminated, truncated, info = self.env.step(action)
        
        # 更新状态
        self.current_observation = observation
        self.current_info = info
        self.episode_reward += reward
        
        # 如果游戏结束，增加回合计数
        if terminated or truncated:
            self.episode_count += 1
            print(f"Episode {self.episode_count}: Reward = {self.episode_reward}, "
                  f"Win = {info.get('player_won', False)}")
        
        return observation, reward, terminated, truncated, info
    
    def run_episode(self, max_steps=1000):
        """运行一个完整的游戏回合"""
        self.reset()
        total_reward = 0
        step_count = 0
        
        while True:
            observation, reward, terminated, truncated, info = self.step()
            total_reward += reward
            step_count += 1
            
            if terminated or truncated or step_count >= max_steps:
                break
        
        return total_reward, info.get('player_won', False), step_count
    
    def close(self):
        """关闭环境"""
        self.env.close()


class TrainedAgent:
    """训练好的AI代理（用于加载预训练模型）"""
    
    def __init__(self, model_path=None):
        self.model = None
        self.action_space = None
        
        if model_path:
            self.load_model(model_path)
    
    def load_model(self, model_path):
        """加载预训练模型"""
        try:
            from stable_baselines3 import PPO
            self.model = PPO.load(model_path)
            print(f"成功加载模型: {model_path}")
        except Exception as e:
            print(f"加载模型失败: {e}")
            print("将使用随机代理")
    
    def select_action(self, observation):
        """根据当前观察选择动作"""
        if self.model is not None:
            # 使用训练好的模型选择动作
            action, _ = self.model.predict(observation, deterministic=True)
            return action
        else:
            # 如果没有模型，使用随机动作
            return random.randint(0, 19)  # 20个可能的动作


# 测试函数
def test_random_agent():
    """测试随机代理"""
    ai = GameAI(render_mode="human")
    
    # 运行几个回合
    for episode in range(3):
        print(f"开始第 {episode+1} 回合")
        reward, won, steps = ai.run_episode()
        print(f"回合结束: 奖励={reward}, 胜利={won}, 步数={steps}")
    
    ai.close()


if __name__ == "__main__":
    test_random_agent()